package com.jsj.subway.ai.agent.rag;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import com.jsj.subway.ai.agent.rag.models.KnowledgeBase;
import com.jsj.subway.ai.largemodel.AliApiClient;
import com.jsj.subway.ai.largemodel.AliClient;
import com.jsj.subway.ai.largemodel.models.AliModelCompletionResponse;
import com.jsj.subway.ai.largemodel.models.Message;
import com.jsj.subway.ai.largemodel.models.Result;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.MultipartBody;
import okhttp3.RequestBody;
import org.apache.commons.lang3.SerializationUtils;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.web.multipart.MultipartFile;

import java.io.*;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;

/**
 * <p>
 * 知识库文件表 服务实现类
 * </p>
 *
 * @author zhouxiaofeng
 * @since 2024-03-14
 */
@Slf4j
@Service
public class KnowledgeBaseFilesServiceImpl  implements KnowledgeBaseFilesService {

    @Autowired
    private KbClient kbClient;

    @Value("${kb.modelFormat}")
    private String modelFormat;

    @Value("${kb.systemPrompt}")
    private String systemPrompt;

    @Value("${kb.isNeedSummarize}")
    private boolean isNeedSummarize;

    @Value("${kb.modelHistorySummarize}")
    private String modelHistorySummarize;

    @Value("${kb.chunkSize}")
    private String chunkSizeDefault;

    @Value("${kb.chunkOverlap}")
    private String chunkOverlapDefault;

    @Value("${kb.fileMaxSize}")
    private Long fileMaxSize;

//    @Autowired
//    @Qualifier("cvmart.queueExecutor")
//    private Executor queueExecutor;

    @Autowired
    private AliClient aliClient;

    public static final String kbCommonName = "common";



    @Override
    public List<Message> queryKbBase(List<Message> messages, boolean isNeedSummarize) {
        if (CollectionUtils.isEmpty(messages)) {
            return messages;
        }
        List<Message> context = new ArrayList<>();
        List<Message> messageNew = SerializationUtils.clone((LinkedList<Message>) messages);
        String query = "";
        String origin = (String) messageNew.get(messageNew.size() - 1).getContent();
        Message mme = messages.get(messageNew.size() - 1);
        try {
            Message systemMsg = new Message();
            systemMsg.setRole("system");
            systemMsg.setContent(systemPrompt);
            context.add(systemMsg);
            List<Message> history = new ArrayList<>();
            if (messageNew.size() > 2) {
                history = messageNew.subList(messageNew.size() - 3, messageNew.size() - 1);
            } else {
                history = messageNew.subList(0, messageNew.size() - 1);
            }
            String history_prompt = "";
            if (history.size() > 0) {
                history_prompt = history.stream().map((x) -> {
                    return (String) x.getContent();
                }).collect(Collectors.joining("\n"));
            }
            String lastContent = StrUtil.format(modelHistorySummarize, history_prompt, mme.getContent());
            mme.setContent(lastContent);
            context.add(mme);
            log.info("lastContent:{}", lastContent);
            if (isNeedSummarize) {
                AliModelCompletionResponse response = aliClient.completionsChat(context);
                if (response.getErr_code() == 0) {
                    query = (String) response.getChoices().get(0).getMessage().getContent();
                    if (query != null && query.indexOf("提问：") == 0) {
                        query = query.substring(3, query.length() - 1);
                    }
                } else {
                    query = origin;
                }
            } else {
                query = origin;
            }
            Result.KbBase kbBase = kbClient.queryKbBase(new KnowledgeBase.QueryReq(kbCommonName, query));
            JSONArray jsonArray = new JSONArray(kbBase.getDocuments());
            log.info("query：{},从文件中召回信息：{}", query, jsonArray);
//            if (botKey == 3) {
//                Result.KbBase kbBase2 = kbClient.queryKbBase(new KnowledgeBase.QueryReq("kbSession" + sessionId, query));
//                JSONArray jsonArray2 = new JSONArray(kbBase2.getDocuments());
//                log.info("文档问答从文件中召回信息：{}", jsonArray2);
//                if (!jsonArray2.isEmpty()) {
//                    jsonArray.addAll(jsonArray2);
//                }
//            }
            if (jsonArray.isEmpty()) {
                return messageNew;
            } else {
                String kn_content = jsonArray.stream().map(x -> x.toString()).collect(Collectors.joining("\n"));
                query = StrUtil.format(modelFormat, kn_content, origin);
            }
            messageNew.get(messages.size() - 1).setContent(query);
            return messageNew;
        } catch (Exception e) {
            log.info("从文件中召回失败， msg：{}", e.getMessage());
            return messageNew;
        }
    }

}
